import os
import sys
sys.path.append('..')
import cv2
import numpy as np
import pandas as pd
from pathlib import Path
import glob
from utils.json import json_to_image

def get_mask(img_path, dst_path, no_defect_code, is_charm, for_t7):
    code = Path(img_path).parent.name

    img = cv2.imread(img_path, 1)


    if is_charm :
        new_img = cv2.resize(img, (int(img.shape[1]*0.517), int(img.shape[0]*0.517)), interpolation=cv2.INTER_LINEAR)
        center_x = int(new_img.shape[1] / 2)
        center_y = int(new_img.shape[0] / 2)
        if for_t7:
            crop_img = new_img[center_y-516:center_y+517, center_x-516:center_x+517]
        else:
            crop_img = new_img[center_y-511:center_y+511, center_x-511:center_x+511]
    else:
        new_img = cv2.resize(img, (int(img.shape[1]*0.361), int(img.shape[0]*0.361)), interpolation=cv2.INTER_LINEAR)
        center_x = int(new_img.shape[1] / 2)
        center_y = int(new_img.shape[0] / 2)
        # print(new_img.shape)
        crop_img = new_img[center_y-487:center_y+487, center_x-605:center_x+606]


    json_path = str(Path(img_path).with_suffix('.json'))
    if (not os.path.exists(json_path)) and  code != no_defect_code:
        return

    try:
        img_name = Path(img_path).name
        img_dst_path = os.path.join(dst_path, code, img_name)
        os.makedirs(os.path.join(dst_path, code), exist_ok=True)
        cv2.imwrite(img_dst_path, crop_img)

        # print(json_path)
        if(os.path.exists(json_path)):
            # print('get mask ...')
            mask = json_to_image(img_shape=img.shape[:2], json_path=json_path)
            mask = mask.astype(np.uint8)
            mask *= 255

            # print('get mask done !')
            if is_charm :
                new_mask = cv2.resize(mask, (int(mask.shape[1]*0.517), int(mask.shape[0]*0.517)), interpolation=cv2.INTER_LINEAR)
                center_x = int(new_mask.shape[1] / 2)
                center_y = int(new_mask.shape[0] / 2)
                if for_t7:
                    crop_mask = new_mask[center_y-516:center_y+517, center_x-516:center_x+517]
                else:
                    crop_mask = new_mask[center_y-511:center_y+511, center_x-511:center_x+511]
            else:
                new_mask = cv2.resize(mask, (int(mask.shape[1]*0.361), int(mask.shape[0]*0.361)), interpolation=cv2.INTER_LINEAR)
                center_x = int(new_mask.shape[1] / 2)
                center_y = int(new_mask.shape[0] / 2)
                crop_mask = new_mask[center_y-487:center_y+487, center_x-605:center_x+606]
            
            mask_name = Path(img_path).with_suffix('.png').name
            mask_dst_path = os.path.join(dst_path, code, mask_name)
            cv2.imwrite(mask_dst_path, crop_mask)
    except:
        return 


def main():
    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6010/all_0426'
    dst_path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6_Vtech_for_Vtech/T6010/0621'
    no_defect_code = "TGXID"
    is_charm = False
    for_t7 = False


    df = pd.DataFrame()
    df['img_path'] = glob.glob(os.path.join(path, "*/*.jpg"))
    print(len(df))

    if is_pandarallel:
        df.parallel_apply(lambda x:get_mask(img_path=x['img_path'], dst_path=dst_path, no_defect_code=no_defect_code, is_charm=is_charm, for_t7=for_t7), axis=1)
    else:
        df.apply(lambda x:get_mask(img_path=x['img_path'], dst_path=dst_path, no_defect_code=no_defect_code, is_charm=is_charm, for_t7=for_t7), axis=1)

    print()


if __name__=='__main__':
    main()
